import backtrader as bt

from indicators.ssa_index_ind import ssa_index_ind


class MyStrategy(bt.Strategy):
    # 策略参数
    params = dict(
        period=20,  # 均线周期
        look_back_days=30,
        printlog=False,
        period_ema_fast=10,
        period_ema_slow=100,
        short_window=15,
        long_window=20,
        ssa_window=15
    )
    def __init__(self):
        for data in self.datas:
           data.mas = dict()
           data.ema_fast = bt.indicators.EMA(period=self.p.period_ema_fast)
           data.ema_slow = bt.indicators.EMA(period=self.p.period_ema_slow)
           data.rsi = bt.indicators.RelativeStrengthIndex()
           data.profit= 0
           data.short_ma = bt.indicators.SMA(data.close, period=self.p.short_window)
           data.long_ma = bt.indicators.SMA(data.close, period=self.p.long_window)
           # 五日移动平均线
           data.sma5 = bt.indicators.SimpleMovingAverage(data, period=5)
           # 十日移动平均线
           data.sma10 = bt.indicators.SimpleMovingAverage(data, period=10)
           data.mas = bt.ind.SMA(data.close, period=self.p.period)
           data.ssa = ssa_index_ind(ssa_window=self.params.ssa_window,data=data, subplot=False)

        #遍历所有股票,计算20日均线
        #for data in self.datas:

    def notify_order(self, order):
        """
        订单状态处理

        Arguments:
            order {object} -- 订单状态
        """
        if order.status in [order.Submitted, order.Accepted]:
            # 如订单已被处理，则不用做任何事情
            return

        # 检查订单是否完成
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                #print('BUY成交, 执行价={0}, {1}'.format(order.executed.price, order.executed.size))
            #elif order.issell():
                #print('SELL成交, 执行价={0}, {1}'.format(order.executed.price, order.executed.size))
            self.bar_executed = len(self)



        # 订单因为缺少资金之类的原因被拒绝执行
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        # 订单状态处理完成，设为空
        self.order = None

    def notify_trade(self, trade):
        """
        交易成果

        Arguments:
            trade {object} -- 交易状态
        """
        if not trade.isclosed:
            return

        # 显示交易的毛利率和净利润
        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
                 (trade.pnl, trade.pnlcomm), doprint=True)



    def next(self):
        #计算截面收益率
        rate_list=[]
        #for data in self.datas:
        #    if len(data)>self.p.look_back_days:
        #        p0=data.close[0]
        #        pn=data.close[-self.p.look_back_days]
        #        rate=(p0-pn)/pn
        #        rate_list.append([data._name,rate])
        #股票池
        long_list=[]
        sorted_rate=sorted(rate_list,key=lambda x:x[1],reverse=True)
        long_list=[i[0] for i in sorted_rate[:10]]
        # 得到当前的账户价值
        total_value = self.broker.getvalue()
        p_value = total_value*0.9/10
        for data in self.datas:
            #获取仓位
            pos = self.getposition(data).size
            #print(len(data.sma5))
            if not pos :
                #print(data.rsi[0])
                #if data.sma5[0] > data.sma10[0] and data.rsi < 30:
                if  data.sma5[0] > data.sma10[0] and data.rsi < 40:
                    #self.order = self.buy()
                    #size=int(p_value/100/data.close[0])*100
                    #self.buy(data = data, size = size)
                    #size = int(p_value / 100 / data.close[0]) * 100
                    self.buy(data=data)
                    #self.order = self.close(exectype=bt.Order.StopTrail, trailpercent=0.1)
            if pos!=0 :
                if  ( (data.close[0]  < data.ssa[0] ) or  data.rsi > 70  ):
                    #self.order = self.sell()
                    #(data.short_ma[-1] > data.long_ma[-1] and data.short_ma[0] < data.long_ma[0] ) or
                    self.close(data = data)
    def log(self, txt, dt=None,doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    #记录交易收益情况（可省略，默认不输出结果）
    #def notify_trade(self,trade):
    #    if not trade.isclosed:
    #        return
    #    self.log(f'策略收益：\n毛收益 {trade.pnl:.2f}, 净收益 {trade.pnlcomm:.2f}')
